package cn.com.zetatech.loader.entity;

import org.apache.parquet.io.SeekableInputStream;

import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * @author jian.chen@zetatech.com.cn
 * @Description TODO
 * 辅助类：用于ParquetReader处理SeekableInputStream
 * @Date 2024/8/15 16:07
 */
public class SeekableByteArrayInputStream extends SeekableInputStream {
    private final byte[] data;
    private int position = 0;

    public SeekableByteArrayInputStream(byte[] data) {
        this.data = data;
    }

    @Override
    public void seek(long newPos) throws IOException {
        if (newPos < 0 || newPos > data.length) {
            throw new IOException("无效的Seek位置: " + newPos);
        }
        position = (int) newPos;
    }

    @Override
    public long getPos() throws IOException {
        return position;
    }

    @Override
    public int read() throws IOException {
        if (position >= data.length) {
            return -1;
        }
        return data[position++] & 0xFF;
    }

    @Override
    public int read(byte[] b, int off, int len) throws IOException {
        if (position >= data.length) {
            return -1;
        }
        int bytesRead = Math.min(len, data.length - position);
        System.arraycopy(data, position, b, off, bytesRead);
        position += bytesRead;
        return bytesRead;
    }

    @Override
    public void readFully(byte[] bytes, int i, int i1) throws IOException {
        read(bytes, i, i1);
    }

    @Override
    public void readFully(byte[] bytes) throws IOException {
        read(bytes, 0, bytes.length);
    }

    @Override
    public int read(ByteBuffer byteBuffer) throws IOException {
        int bytesToRead = Math.min(byteBuffer.remaining(), data.length - position);
        byteBuffer.put(data, position, bytesToRead);
        position += bytesToRead;
        return bytesToRead;
    }

    @Override
    public void readFully(ByteBuffer byteBuffer) throws IOException {
        read(byteBuffer);
    }

    @Override
    public void close() throws IOException {
        // 内存字节数组无需关闭
    }
}
